#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch
from torch import Tensor
import torch.nn as nn
import copy
import math


def clones(obj, n):
    return nn.ModuleList([copy.deepcopy(obj) for _ in range(n)])


class LayerNorm(nn.Module):
    def __init__(self):
        super(LayerNorm, self).__init__()

    def forward(self):
        pass


class ScaledDotAttention(nn.Module):
    """
    Q: query   (Batch, Q_len, Q_dim)
    K: key     (Batch, K_len, K_dim)
    V: Value   (Batch, V_len, V_dim)
    Q_dim = K_dim = V_dim


    """
    def __init__(self, d_model, dropout):
        super(ScaledDotAttention, self).__init__()
        self.scale = math.sqrt(d_model)
        self.softmax = nn.Softmax(dim=2)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask=None):
        attn = torch.bmm(query, key.transpose(1, 2))  # (Batch, Q_len, K_len)

        attn *= self.scale

        if mask:
            attn = attn.masked_fill(mask, 1e-20)
        attn = self.softmax(attn)
        attn = self.dropout(attn)
        context = torch.bmm(attn, value)
        return context, attn


class MultiHeadAttention(nn.Module):
    def __init__(self, num_heads):
        super(MultiHeadAttention, self).__init__()

    def forward(self):
        pass








class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()

    def forward(self, *input):
        pass


class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()

    def forward(self, *input):
        pass


class AddNorm(nn.Module):
    def __init__(self):
        super(AddNorm, self).__init__()

    def forward(self, *input):
        pass


class PositionWiseFeedForward(nn.Module):
    def __init__(self):
        super(PositionWiseFeedForward, self).__init__()

    def forward(self, *input):
        pass


class MultiHeadAttention(nn.Module):
    def __init__(self):
        super(MultiHeadAttention, self).__init__()

    def forward(self, *input):
        pass


class Embedding(nn.Module):
    def __init__(self, vocab_size, d_model):
        super(Embedding, self).__init__()
        self.emb = nn.Embedding(vocab_size, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.emb(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    """
    PE(pos, 2*i) = sin(pos/10000^(2*i/d_model))
    PE(pos, 2*i+1) = cos(pos/10000^(2*i/d_model))
    在序列每个位置上，对该位置做不同弧度的编码，偶数位置做sin，奇数位置做cos

    """
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        pe = torch.zeros(size=(max_len, d_model), requires_grad=False)  # (max_len, d_model)
        positions = torch.arange(0., max_len, 1.).unsqueeze(1)  # (max_len, 1)

        div_term = torch.exp(torch.arange(0., d_model, 2.) * -(math.log(10000.0) * d_model)) # (
        pe[:, 0::2] = torch.sin(positions * div_term)
        print(pe[:, 0::2].shape)
        pe[:, 1::2] = torch.cos(positions * div_term)
        pe = pe.unsqueeze(0)  # (1, max_len， d_model)

        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)

# nn.LayerNorm


def test():
    pass


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import numpy as np
    plt.figure(figsize=(15, 5))
    pe = PositionalEncoding(20, 0)
    y = pe.forward(torch.zeros(1, 100, 20))
    plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
    plt.legend(["dim %d" % p for p in [4, 5, 6, 7]])
    plt.show()
